from cookielib import CookieJar
import os
from urllib import urlencode
from urllib2 import build_opener, HTTPCookieProcessor
import sys


class HDOJ():

    def __init__(self, username='', userpass='', hdoj_url='http://acm.hdu.edu.cn/'):

        self.username = username
        self.userpass = userpass
        self.cookiejar = CookieJar()
        self.opener = build_opener(HTTPCookieProcessor(self.cookiejar))
        self.opener.addheaders = [('User-agent', 'Mozilla/4.0 (compatible; MSIE 6.0; Windows NT 5.1)')]
        self.hdoj_url = hdoj_url
        self.solve = []

    def login(self):

        body = urlencode({'username': self.username, 'userpass': self.userpass, 'login': 'Sign In'})
        self.opener.open(self.hdoj_url + '/userloginex.php?action=login', body)

    def get_solved(self):

        op = self.opener.open(self.hdoj_url + 'userstatus.php?user=' + self.username)
        data = str(op.read())
        p = data.split('List of solved problems')[1].split('<script language=javascript>')[1].split(';</script>')[0]
        pid_ac_sub = p.split('p(')[1:]
        self.solve = [i.split(',')[0] for i in pid_ac_sub]

    def save_source_code(self):

        if not os.path.exists('src'):
            os.mkdir('src')
        ext = {'G++': '.cpp',
               'GCC': '.c',
               'C++': '.cpp',
               'C': '.c',
               'Pascal': '.pas',
               'Java': '.java'}
        for i in self.solve:
            if os.path.exists(os.path.join('src', '%s.cpp' % i)) or os.path.exists(os.path.join('src', '%s.c' % i))\
                    or os.path.exists(os.path.join('src', '%s.pas' % i)) or os.path.exists(os.path.join('src', '%s.java' % i)):
                continue
            op = self.opener.open(self.hdoj_url + 'status.php?user=%s&pid=%s&status=5' % (self.username, i))
            data = str(op.read())
            rid = data.split('/viewcode.php?rid=')[1].split('"')[0]
            op = self.opener.open(self.hdoj_url + 'viewcode.php?rid=%s' % rid)
            data = str(op.read())
            language = data.split('Language : ')[1].split('&')[0]
            source_code = data.split('<textarea id=usercode style="display:none;text-align:left;">')[1].split('</textarea>')[0]
            source_code = source_code.replace('&nbsp;', ' ')
            source_code = source_code.replace('&lt;', '<')
            source_code = source_code.replace('&gt;', '>')
            source_code = source_code.replace('&amp;', '&')
            source_code = source_code.replace('&quot;', '"')
            source_code = source_code.replace('&apos;', "'")
            f = open(os.path.join('src', '%s%s' % (i, ext[language])), 'w')
            f.write(source_code)
            f.close()
            print 'get %s' % i


def main():
    username = sys.argv[1]
    userpass = sys.argv[2]
    hdoj = HDOJ(username=username, userpass=userpass)
    hdoj.login()
    hdoj.get_solved()
    hdoj.save_source_code()

if __name__ == '__main__':
    main()

